﻿#pragma once
enum Colour
{
	RED,
	BLACK
};

template<class T>
struct RBTreeNode
{
	T _data;
	RBTreeNode<T>* _left;
	RBTreeNode<T>* _right;
	RBTreeNode<T>* _parent;
	Colour _col;

	RBTreeNode(const T& data)
		:_left(nullptr)
		,_right(nullptr)
		,_parent(nullptr)
		,_data(data)
	{}

};

template<class T,class Ref,class Ptr>
struct RBTreeIterator
{
	typedef RBTreeNode<T> Node;
	typedef RBTreeIterator<T, Ref, Ptr> Self;

	Node* _node;
	Node* _root;

	RBTreeIterator(Node* node,Node* root)
		:_node(node)
		,_root(root)
	{}

public:
	Self& operator++()
	{
		if (_node->_right)
		{
			//_node的左孩子存在
			Node* cur = _node->_right;
			while (cur->_left)
			{
				cur = cur->_left;
			}
			_node = cur;
		}
		else
		{
			Node* cur = _node;
			Node* parent = cur->_parent;

			while (parent && cur == parent->_right)
			{
				cur = parent;
				parent = cur->_parent;
			}

			_node = parent;
		}

		return *this;
	}

	Self& operator--()
	{
		if (_node == nullptr) // end()
		{
			// --end()，特殊处理，⾛到中序最后⼀个结点，整棵树的最右结点
			Node* rightMost = _root;
			while (rightMost && rightMost->_right)
			{
				rightMost = rightMost->_right;
			}
			_node = rightMost;
		}
		else if (_node->_left)
		{
			// 左⼦树不为空，中序左⼦树最后⼀个
			Node* rightMost = _node->_left;
			while (rightMost->_right)
			{
				rightMost = rightMost->_right;
			}
			_node = rightMost;
		}
		else
		{
			// 孩⼦是⽗亲右的那个祖先
			Node* cur = _node;
			Node* parent = cur->_parent;
			while (parent && cur == parent->_left)
			{
				cur = parent;
				parent = cur->_parent;
			}
			_node = parent;
		}
		return *this;
	}
	Ref operator*()
	{
		return _node->_data;
	}
	Ptr operator->()
	{
		return &_node->_data;
	}
	bool operator!= (const Self& s) const
	{
		return _node != s._node;
	}
	bool operator== (const Self& s) const
	{
		return _node == s._node;
	}
};

template<class K,class T,class _Compare>
class RBTree
{
	typedef RBTreeNode<T> Node;
public:
	typedef RBTreeIterator<T, T&, T*> Iterator;
	typedef RBTreeIterator<T, const T&, const T*> ConstIterator;

	Iterator Begin()
	{
		Node* leftMost = _root;
		while (leftMost && leftMost->_left)
		{
			leftMost = leftMost->_left;
		}
		return Iterator(leftMost, _root);//构造一个迭代	器返回
	}

	Iterator End()
	{
		return Iterator(nullptr, _root);
	}
	ConstIterator Begin() const
	{
		Node* leftMost = _root;
		while (leftMost && leftMost->_left)
		{
			leftMost = leftMost->_left;
		}
		return ConstIterator(leftMost, _root);
	}
	ConstIterator End() const
	{
		return ConstIterator(nullptr, _root);
	}
	RBTree() = default;
	//插入
	bool Insert(const T& data)
	{
		if (_root == nullptr)
		{
			_root = new Node(data);
			_root->_col = BLACK;
			return true;
		}

		Node* cur = _root;
		Node* parent = nullptr;
		_Compare com;
		while (cur)
		{
			if (com(data) < com(cur->_data))
			{
				parent = cur;
				cur = cur->_left;
			}
			else if (com(data) > com(cur->_data))
			{
				parent = cur;
				cur = cur->_right;
			}
			else
				return false;
		}

		cur = new Node(data);
		cur->_col = RED;

		if (com(cur->_data) < com(parent->_data) )
		{
			parent->_left = cur;
			cur->_parent = parent;
		}
		else
		{
			parent->_right = cur;
			cur->_parent = parent;
		}

		
		while (parent && parent->_col == RED)
		{
			Node* grandfather = parent->_parent;
			
			if (parent == grandfather->_left)
			{
				//    g
				//  p   u
				// c
				//情况1：只变色
				Node* uncle = grandfather->_right;
				if (uncle && uncle->_col == RED)
				{
					//变色
					parent->_col = BLACK;
					uncle->_col = BLACK;
					grandfather->_col = RED;

					//
					parent = grandfather->_parent;
					cur = grandfather;
				}
				else
				{
					//情况2
					//      g
					//   p     u(不存在或者为黑）
					// c
					//1:cur是左孩子
					if (cur == parent->_left)
					{
						//

						RotateR(grandfather);

						parent->_col = BLACK;
						grandfather->_col = RED;
					}
					else
					{
						RotateL(parent);
						RotateR(grandfather);
						cur->_col = BLACK;
						grandfather->_col = RED;
					}
					break;
				}
			}
			else
			{
				Node* uncle = grandfather->_left;
				if (uncle && uncle->_col == RED)
				{
					//变色
					parent->_col = BLACK;
					uncle->_col = BLACK;
					grandfather->_col = RED;

					//
					parent = grandfather->_parent;
					cur = grandfather;

				}
				else
				{
					if (cur == parent->_right)
					{
						RotateL(grandfather);

						grandfather->_col = RED;
						parent->_col = BLACK;
						//cur->_col = BLACK;

					}
					else
					{
						RotateR(parent);
						RotateL(grandfather);
						grandfather->_col = RED;
						//parent->_col = BLACK;
						cur->_col = BLACK;
					}
					break;
				}
			}
		}
		_root->_col = BLACK;
		return true;
	}

	void RotateR(Node* parent)
	{
		Node* praL = parent->_left;
		Node* praLR = praL->_right;

		parent->_left = praLR;
		praL->_right = parent;

		Node* grandfather = parent->_parent;

		if (parent == _root)
		{
			praL->_parent = nullptr;
			_root = praL;
		}
		else
		{
			if (parent == grandfather->_left)
				grandfather->_left = praL;
			else
				grandfather->_right = praL;

			praL->_parent = grandfather;
		}
		parent->_parent = praL;
		if (praLR)
			praLR->_parent = parent;

	}
	/*void RotateR(Node* parent)
	{
		Node* subL = parent->_left;
		Node* subLR = subL->_right;

		parent->_left = subLR;
		if (subLR)
			subLR->_parent = parent;

		Node* pParent = parent->_parent;

		subL->_right = parent;
		parent->_parent = subL;

		if (parent == _root)
		{
			_root = subL;
			subL->_parent = nullptr;
		}
		else
		{
			if (pParent->_left == parent)
			{
				pParent->_left = subL;
			}
			else
			{
				pParent->_right = subL;
			}

			subL->_parent = pParent;
		}
	}*/
	void RotateL(Node* parent)
	{
		Node* subR = parent->_right;
		Node* subRL = parent->_right->_left;

		parent->_right = subRL;
		if (subRL)
			subRL->_parent = parent;

		subR->_left = parent;

		Node* pParent = parent->_parent;

		parent->_parent = subR;
		subR->_parent = pParent;

		if (pParent == nullptr)
		{
			_root = subR;
			subR->_parent = nullptr;;

		}
		else
		{
			if (pParent->_left == parent)
			{
				pParent->_left = subR;
			}
			else
			{
				pParent->_right = subR;

			}
			subR->_parent = pParent;

		}
	}
	int Height()
	{
		return _Height(_root);
	}

private:
	int _Height(Node * root)
	{
		if (root == nullptr)
			return 0;
		int leftHeight = _Height(root->_left);
		int rightHeight = _Height(root->_right);
		return leftHeight > rightHeight ? leftHeight + 1 : rightHeight + 1;
	}
private:
	Node* _root = nullptr;
};